# 1. 导入依赖（新增工具相关）
from typing import Annotated, Optional, Literal
from typing_extensions import TypedDict
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from langchain.chat_models import init_chat_model
from langchain.schema import AIMessage, SystemMessage

import os




# 3. 定义 State（新增工具相关字段）
class ToolChatState(TypedDict):
    messages: Annotated[list, add_messages]  # 对话历史
    need_tool: Optional[bool] = None  # 是否需要调用工具（初始为 None）
    tool_result: Optional[str] = None  # 工具调用结果（初始为 None）


# 4. 初始化 LLM
llm = init_chat_model(
    model="gpt-4o-mini",
    model_provider="openai",
    api_key="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",
    temperature=0.3
)


# 5. 定义节点 1：判断是否需要调用工具（judge_tool 节点）
def judge_tool_node(state: ToolChatState) -> dict:
    """判断用户问题是否需要调用工具（仅天气相关问题需要）"""
    # 构造判断提示词
    judge_prompt = SystemMessage(
        content="""你是一个工具调用判断器。规则如下：
        1. 如果用户问题是查询天气（如“北京今天天气”“上海明天是否下雨”），则需要调用工具，返回 need_tool=True
        2. 其他问题（如常识、聊天、知识问答）不需要调用工具，返回 need_tool=False
        仅需返回 {"need_tool": True/False}，不要附加其他内容！"""
    )
    # 传入用户问题和判断提示词
    input_messages = [judge_prompt] + state["messages"]
    # 调用 LLM 判断
    judge_result = llm.invoke(input_messages)
    # 解析判断结果（提取 need_tool 布尔值）
    import json
    try:
        need_tool = json.loads(judge_result.content)["need_tool"]
    except:
        need_tool = False  # 解析失败时默认不调用工具
    # 返回更新的 State 字段
    return {"need_tool": need_tool}


# 6. 定义节点 2：调用工具（call_tool 节点，模拟天气查询）
def call_tool_node(state: ToolChatState) -> dict:
    """模拟调用天气工具（实际项目中可替换为真实 API，如高德天气、OpenWeather）"""
    # 从对话历史中提取用户问题（获取最新的用户消息）
    user_message = next(
        msg for msg in reversed(state["messages"])
        if msg.type == "human"
    ).content

    # 模拟工具调用（实际场景中替换为 requests 调用真实 API）
    print(f"[工具调用中] 处理用户问题：{user_message}")
    # 模拟返回天气结果（实际应从 API 获取）
    tool_result = f"【天气查询结果】根据模拟数据，{user_message.split('，')[0] if '，' in user_message else user_message}：今日晴，气温 25-32℃，微风。"

    # 返回更新的 State 字段（工具结果）
    return {"tool_result": tool_result}


# 7. 定义节点 3：生成最终回答（chatbot 节点）
def chatbot_node(state: ToolChatState) -> dict:
    """根据是否有工具结果，生成最终回答"""
    # 构造提示词（包含工具结果（如有）和对话历史）
    system_prompt = SystemMessage(
        content=f"""你是一个智能助手，根据以下信息生成回答：
        - 如果有工具结果（tool_result），必须基于工具结果回答，不要编造信息
        - 如果没有工具结果，直接基于常识回答
        工具结果（如有）：{state.get('tool_result', '无')}"""
    )
    # 传入提示词和对话历史
    input_messages = [system_prompt] + state["messages"]
    # 调用 LLM 生成最终回答
    ai_response = llm.invoke(input_messages)
    # 返回更新的 State 字段（追加 AI 回答）
    return {"messages": [ai_response]}


# 8. 定义条件分支函数：决定从 judge_tool 节点跳转到哪个节点
def tool_branch(state: ToolChatState) -> Literal["call_tool", "chatbot"]:
    """条件分支逻辑：根据 need_tool 决定跳转方向"""
    return "call_tool" if state["need_tool"] else "chatbot"


# 9. 构建 Graph（图）
graph_builder = StateGraph(ToolChatState)

# 9.1 添加所有节点
graph_builder.add_node("judge_tool", judge_tool_node)  # 判断工具节点
graph_builder.add_node("call_tool", call_tool_node)  # 调用工具节点
graph_builder.add_node("chatbot", chatbot_node)  # 生成回答节点

# 9.2 添加边（包含条件边）
graph_builder.add_edge(START, "judge_tool")  # 开始 → 判断节点
# 条件边：从 judge_tool 节点根据 tool_branch 结果跳转
graph_builder.add_conditional_edges(
    source="judge_tool",  # 源节点
    path=tool_branch,  # 分支判断函数
    # 分支映射：判断结果 → 目标节点
    path_map={
        "call_tool": "call_tool",
        "chatbot": "chatbot"
    }
)
graph_builder.add_edge("call_tool", "chatbot")  # 工具节点 → 回答节点
graph_builder.add_edge("chatbot", END)  # 回答节点 → 结束

# 9.3 编译图
tool_chat_graph = graph_builder.compile()

# 10. 可视化图结构
try:
    from IPython.display import Image
    png_data = tool_chat_graph.get_graph().draw_mermaid_png()
    with open("graph.png", "wb") as f:
        f.write(png_data)
    print("图表已保存为 graph.png")
except Exception as e:
    print(f"保存图表失败：{e}")


# 11. 运行交互逻辑
def run_tool_chatbot():
    print("=== 带工具调用判断的 LangGraph 机器人 ===")
    print("提示：输入空字符串退出，可尝试问“北京今天天气”或“什么是 AI”")

    while True:
        user_input = input("\nUser: ")
        if user_input.strip() == "":
            print("对话结束！")
            break

        # 初始化 State：传入用户输入
        initial_state = {"messages": [{"role": "user", "content": user_input}]}

        # 流式执行图并输出结果
        print("Assistant: ", end="")
        for event in tool_chat_graph.stream(initial_state):
            for node_output in event.values():
                if "messages" in node_output:
                    latest_msg = node_output["messages"][-1]
                    if isinstance(latest_msg, AIMessage):
                        print(latest_msg.content, end="")
    print("\n")


# 12. 启动机器人
if __name__ == "__main__":
    run_tool_chatbot()